package arena

import (
	"fmt"
	"sync"
)

type ArenaLock struct {
	mu     sync.Mutex
	locks  map[string]*sync.Mutex
	preKey string
}

func NewArenaLock(preKey string) *ArenaLock {
	return &ArenaLock{
		locks:  make(map[string]*sync.Mutex),
		preKey: preKey,
	}
}

func (al *ArenaLock) lockKey(key string) {
	al.mu.Lock()
	if _, exists := al.locks[key]; !exists {
		al.locks[key] = &sync.Mutex{}
	}
	lock := al.locks[key]
	al.mu.Unlock()

	lock.Lock()
}

func (al *ArenaLock) unlockKey(key string) {
	al.mu.Lock()
	if lock, exists := al.locks[key]; exists {
		lock.Unlock()
		delete(al.locks, key)
	}
	al.mu.Unlock()
}

func (al *ArenaLock) Lock(keys ...string) error {
	if len(keys) == 0 {
		return fmt.Errorf("no keys provided")
	}

	for i := range keys {
		keys[i] = al.preKey + keys[i]
	}

	// Ensure keys are locked in order to prevent deadlocks
	for _, key := range keys {
		al.lockKey(key)
	}

	return nil
}

func (al *ArenaLock) Unlock(keys ...string) {
	if len(keys) == 0 {
		return
	}

	for i := range keys {
		keys[i] = al.preKey + keys[i]
	}

	// Unlock in reverse order
	for i := len(keys) - 1; i >= 0; i-- {
		al.unlockKey(keys[i])
	}
}
